import { Context, h, Schema } from 'koishi'
import { Buffer } from 'buffer'

export const inject = {
  required: ['http'],
}

export const name = 'awa-image-recognizer'

export interface Config {
  ollamaUrl: string
  imageRecognitionModel: string
  geminiUrl: string
  defaultLanguage: string
}

export const Config: Schema<Config> = Schema.object({
  ollamaUrl: Schema.string().default('http://127.0.0.1:11434').description('Ollama 后端地址'),
  imageRecognitionModel: Schema.string().default('moondream').description('用于图片描述的模型名称，例如 moondream'),
  geminiUrl: Schema.string().default('http://192.168.31.84:8629').description('Gemini 后端地址'),
  defaultLanguage: Schema.string().default('en_us').description('默认语言，zh_cn 表示中文，en 表示英文'),
})


export function apply(ctx: Context, config: Config) {
  ctx.command('air')
    .option("language", "-l, --language <language:string>")
    .alias('awa_image_recognizer')
    .action(async ({ session, options }) => {
      const LANGUAGE = options.language ? options.language : config.defaultLanguage;

      // 1. 检查是否回复了消息，并且消息中包含图片
      if (!session.quote) {
        await session.send(`${h.quote(session.messageId)}请回复一张图片。`)
        return
      }

      const elements = h.parse(session.quote.content)
      const imageElement = elements.find(e => e.type === 'img')

      if (!imageElement) {
        await session.send(`${h.quote(session.messageId)}回复的消息中没有图片。`)
        return
      }

      const imageUrl = imageElement.attrs.src
      ctx.logger.info(`Found image URL: ${imageUrl}`)

      if (!imageUrl) {
        await session.send(`${h.quote(session.messageId)}图片 URL 不存在。`)
        return
      }

      // 2. 下载图片并转换为 Base64
      try {
        const hint1MsgId = await session.send(`${h.quote(session.messageId)}正在处理图片，请稍等...`)

        // 使用 Koishi 的 http 服务下载图片，直接获取 ArrayBuffer
        const imageData = await ctx.http.get<ArrayBuffer>(imageUrl, { responseType: 'arraybuffer' })
        const base64Image = Buffer.from(imageData).toString('base64')

        await session.bot.deleteMessage(session.guildId, String(hint1MsgId));

        // 3. 调用图片描述模型获取英文描述
        const hint2MsgId = await session.send(`${h.quote(session.messageId)}正在调用图片识别模型...`)
        const imageRecognitionBody = {
          model: config.imageRecognitionModel,
          prompt: 'Describe this image in detail.',
          images: [base64Image],
          stream: false,
        }
        const recognitionResult = await ctx.http.post<any>(`${config.ollamaUrl}/api/generate`, imageRecognitionBody)

        let finalResponse = recognitionResult?.response;
        const originalEnglishResponse = recognitionResult?.response;

        await session.bot.deleteMessage(session.guildId, String(hint2MsgId));

        // 4. 如果用户指定了中文，则调用 Gemini 后端进行翻译
        if (LANGUAGE === 'zh_cn' && originalEnglishResponse) {
          const hint3MsgId = await session.send(`${h.quote(session.messageId)}识别完成，正在进行翻译...`)

          // 构建 Gemini 翻译请求体
          const geminiRequestBody = {
            text: `请将以下英文文本精确地、无任何多余内容地翻译成中文，只返回翻译结果，不要返回原始文本：\n\n${originalEnglishResponse}`
          }

          try {
            const geminiResult = await ctx.http.post<any>(`${config.geminiUrl}/gemini`, geminiRequestBody)

            if (geminiResult?.answer) {
              finalResponse = geminiResult.answer.trim();
            } else {
              // Gemini 返回的 JSON 中没有 answer 字段，视为失败
              throw new Error('Gemini response has no answer field');
            }
          } catch (translationError) {
            ctx.logger.warn('Translation failed:', translationError);
            // 翻译失败，使用英文原文并添加提示
            finalResponse = `翻译失败，英文原文：\n${originalEnglishResponse}`;
          }

          await session.bot.deleteMessage(session.guildId, String(hint3MsgId));
        }

        // 5. 返回最终结果
        if (finalResponse) {
          await session.send(`${h.quote(session.messageId)}${finalResponse.trim()}`)
        } else {
          await session.send(`${h.quote(session.messageId)}未能获取到图片描述，可能是 Ollama 或 Gemini 服务出现问题。`)
        }

      } catch (error) {
        ctx.logger.error('Error during image recognition or translation:', error)
        await session.send(`${h.quote(session.messageId)}识别图片时发生错误，请稍后重试。`)
      }
    })
}